import time

from langchain_core.runnables import RunnableLambda
from langchain_core.tracers.schemas import Run
from langchain_core.runnables import RunnableConfig

def on_start(run_obj: Run, config: RunnableConfig) -> None:
    print("on_start")
    print("run_obj: ", run_obj)
    print("config: ", config)
    print("=======================")

def on_end(run_obj: Run, config: RunnableConfig) -> None:
    print("on_end")
    print("run_obj: ", run_obj)
    print("config: ", config)
    print("-----------------------")

def on_error(run_obj: Run, config: RunnableConfig) -> None:
    print("on_error")
    print("run_obj: ", run_obj)
    print("config: ", config)
    print("+++++++++++++++++++++++")


# 1. 创建RunnableLamdba与链
runnable = RunnableLambda(lambda x: time.sleep(x)).with_listeners(
    on_start=on_start,
    on_end=on_end,
    on_error=on_error,
)
chain = runnable

# 2. 调用并执行
content = chain.invoke(2, config={"configuration": {"name": "慕小课"}})

# print(content)
